#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
@Author          : ricsy
@File            : web_asset_sniffer.py
@Create Time     : 2025/6/4 00:01
@Last Modified   :
@Description     :
"""
import codecs
import json
import logging
import os
import time
from pathlib import Path
from urllib.parse import unquote, urlparse

import requests
import ruamel.yaml
from selenium import webdriver
from selenium.webdriver.chrome.options import Options
from selenium.webdriver.chrome.service import Service as ChromeService
from webdriver_manager.chrome import ChromeDriverManager


logging.basicConfig(
    level=logging.INFO,
    format="[%(asctime)s] [%(levelname)s] %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
    encoding="utf-8",
    handlers=[
        logging.StreamHandler(),
    ],
)

yaml_file = "data/webstack.yml"


def get_website_links():
    """读取网址列表"""
    yaml = ruamel.yaml.YAML()
    with codecs.open(yaml_file, "r", encoding="utf-8") as ReadYaml:
        data = yaml.load(ReadYaml)

    website_links = {
        link["url"]
        for taxonomy in data
        if "list" in taxonomy
        for _list in taxonomy["list"]
        if "links" in _list
        for link in _list["links"]
        if link.get("url")
    }

    return list(website_links)


def download_icon(url, save_path):
    """下载图标文件"""
    try:
        response = requests.get(url, stream=True, timeout=10)
        if response.status_code == 200:
            with open(save_path, "wb") as f:
                for chunk in response.iter_content(1024):
                    f.write(chunk)
            return True
    except Exception as e:
        logging.warning(f"下载失败 {url}: {str(e)}")
    return False


def check_url_exists(driver, url):
    """综合验证URL是否有效（支持图片/文件检测）"""
    try:
        # 方法1：直接HEAD请求（最快）
        response = requests.head(url, timeout=5, allow_redirects=True)
        if response.status_code < 400:  # 2xx/3xx状态码
            return True

        # 方法2：Selenium页面内容检测（更准确）
        driver.get(url)
        time.sleep(1)  # 等待加载

        # 多条件验证
        if (
            "404" not in driver.title
            and "Not Found" not in driver.page_source
            and not any(img in driver.page_source for img in ["<svg", "<img"])
            # 非HTML文档
        ):
            return True

    except Exception:
        pass
    return False


def extract_resources(driver, url):
    """提取并下载页面图标资源"""
    driver.get(url)
    time.sleep(2)  # 等待页面加载

    # 阶段1：排除不可能的资源 + 匹配ICON_RELS
    declared_icons = driver.execute_script(
        """
        // 必须排除的资源类型（非图标）
        const EXCLUDED_RELS = [
            'stylesheet',               // 链接到 CSS 样式表，纯文本样式文件，与视觉图标无关

            'preload',                  // 资源预加载，通常用于字体、脚本等非图标内容
            'prefetch',                 // 资源预加载，通常用于字体、脚本等非图标内容
            'modulepreload',            // 资源预加载，通常用于字体、脚本等非图标内容

            'canonical',                // 指定网页的规范URL（SEO用途），仅用于搜索引擎优化，不链接到任何资源文件
            'manifest',                 // PWA（渐进式Web应用）的清单文件，链接的是JSON配置文件，包含应用元数据而非图像资源
        ];

        // 目标图标类型（按优先级排序）
        const ICON_RELS = [
            'apple-touch-icon',         // iOS 设备首选
            'icon',                     // 标准 favicon
            'shortcut icon',            // 旧版 IE
            'mask-icon',                // Safari 固定标签页图标
            'fluid-icon'                // 微软 Fluid 图标
        ];

        return Array.from(document.querySelectorAll('link'))
            .filter(link => {
                const rel = (link.rel || '').toLowerCase();
                const href = link.href;

                // 硬性排除条件
                if (!href || href.startsWith('data:')) return false;
                if (EXCLUDED_RELS.some(ex => rel.includes(ex))) return false;

                // 匹配图标类型
                return ICON_RELS.some(iconRel => rel.includes(iconRel));
            })
            .map(link => ({
                rel: link.rel,
                href: new URL(link.href, window.location.href).href, // 转为绝对路径
                sizes: link.getAttribute('sizes') || 'any',
                type: link.getAttribute('type') || 'unknown',
                source: 'link-tag'
            }));
    """
    )

    # 阶段2：补充<meta>标签声明的图标（如Windows磁贴）
    meta_icons = driver.execute_script(
        """
        const metaIcons = [
            { name: 'msapplication-TileImage', rel: 'mstile' },
            { name: 'msapplication-square150x150logo', rel: 'mstile-150' }
        ];

        return metaIcons.map(meta => {
            const el = document.querySelector(`meta[name="${meta.name}"]`);
            return el ? {
                rel: meta.rel,
                href: new URL(el.content, window.location.href).href,
                sizes: meta.name.includes('x') ? meta.name.split('logo')[0] : 'any',
                type: 'unknown',
                source: 'meta-tag'
            } : null;
        }).filter(Boolean);
    """
    )

    # 阶段3：检测默认路径图标（即使未被声明）
    default_icons = []
    default_paths = [
        "/favicon.ico",
        "/favicon.png",
        "/apple-touch-icon.png",
        "/icon.png",
        "/favicon-32x32.png",
        "/android-chrome-192x192.png",
    ]

    for path in default_paths:
        try:
            full_url = url.rstrip("/") + path
            driver.get(full_url)
            # 判断页面是否存在
            if check_url_exists(driver, full_url):
                default_icons.append(
                    {
                        "rel": "default-" + path.split(".")[-1],
                        "href": full_url,
                        "sizes": extract_size_from_path(path),  # 自定义函数提取尺寸
                        "type": "image/" + path.split(".")[-1],
                        "source": "default-path",
                    }
                )
        except:
            continue

    # 阶段4：合并 + 去重 + 优先级排序
    all_icons = {}
    for icon in [*declared_icons, *meta_icons, *default_icons]:
        if icon["href"] not in all_icons:  # 去重
            all_icons[icon["href"]] = icon

    # 排序规则：优先 apple > 尺寸大 > 标准 icon > 其他
    return sorted(
        all_icons.values(),
        key=lambda x: (
            -x["rel"].count("apple"),  # apple-touch-icon 优先
            (
                -int(x["sizes"].split("x")[0])
                if x["sizes"].replace("x", "").isdigit()
                else 0
            ),
            x["rel"].count("icon"),  # 标准icon次之
            -["link-tag", "meta-tag", "default-path"].index(
                x["source"]
            ),  # 声明优于默认
        ),
    )


def extract_size_from_path(path):
    """从路径中提取尺寸（如 /favicon-32x32.png -> 32x32）"""
    import re

    match = re.search(r"(\d+x\d+)", path)
    return match.group(1) if match else "any"


def save_icons(index, base_url, icons):
    """保存图标到本地"""
    domain = urlparse(base_url).netloc.replace("www.", "")
    output_dir = Path("output/icons")
    os.makedirs(output_dir, exist_ok=True)

    results = []
    for idx, icon in enumerate(icons):
        try:
            # 提取源文件名
            src_filename = unquote(urlparse(icon["href"]).path.split("/")[-1])
            if not src_filename:
                ext = icon["type"].split("/")[-1] if "/" in icon["type"] else "bin"
                src_filename = f"icon_{idx}.{ext}"

            # 生成保存路径
            save_name = f"{index}-{idx}-{domain}-{src_filename}"
            save_path = output_dir / save_name

            # 下载文件
            if download_icon(icon["href"], save_path):
                results.append(
                    {
                        "source_url": base_url,
                        "icon_url": icon["href"],
                        "local_path": str(save_path),
                        "type": icon["type"],
                        "rel": icon["rel"],
                    }
                )
                logging.info(f"Downloaded: {save_path}")
            else:
                logging.warning(f"Failed to download: {icon['href']}")

        except Exception as e:
            logging.error(f"Error processing {icon['href']}: {str(e)}")

    return results


def save_urls(all_data):
    """
    清空并批量保存所有网站的图标数据到单个JSON文件
    :param all_data: 列表格式，包含所有网站的图标数据
    """
    os.makedirs("output", exist_ok=True)
    filename = "output/urls.json"

    # 清空文件并写入所有数据
    with open(filename, "w", encoding="utf-8") as f:
        json.dump(all_data, f, indent=2, ensure_ascii=False)

    logging.info(f"Saved all url data to: {filename}")


def init_driver():
    """初始化浏览器驱动"""
    chrome_options = Options()
    chrome_options.add_argument("--headless")  # 无头模式
    chrome_options.add_argument("--disable-gpu")
    chrome_options.add_argument("--no-sandbox")
    chrome_options.add_argument("--disable-software-rasterizer")

    driver = webdriver.Chrome(
        service=ChromeService(ChromeDriverManager().install()), options=chrome_options
    )

    # 设置超时
    driver.set_page_load_timeout(15)
    return driver


def main():
    """主执行逻辑"""
    websites = get_website_links()
    driver = init_driver()
    all_icons_data = []  # 存储所有网站的图标数据
    no_icon_urls = []  # 存储无法获取图标URL的网站

    try:
        for idx, url in enumerate(websites):
            if not url.startswith(("http://", "https://")):
                url = "https://" + url
            try:
                logging.info(f"[download] [{idx + 1}/{len(websites)}] fetching: {url}")
                icons = extract_resources(driver, url)
                saved_files = save_icons(idx, url, icons)
                if len(saved_files) == 0:
                    no_icon_urls.append(url)
                logging.info(f"Saved {len(saved_files)} icons from {url}\n")

                # 收集数据（不立即写入文件）
                all_icons_data.append(
                    {"website": url, "icons": icons, "saved_files": saved_files}
                )

            except Exception as e:
                print(f"处理 {url} 时出错: {str(e)}")
                continue
        if no_icon_urls:
            logging.warning(
                f"Failed to fetch icons from the following URLs：\n{','.join(no_icon_urls)}"
            )
        # 所有网站处理完成后，一次性写入文件
        save_urls(all_icons_data)
    finally:
        driver.quit()


if __name__ == "__main__":
    main()
